package cn.seecoder.ai.algorithm.ml.classification

import cn.seecoder.ai.algorithm.ml.BaseClassification
import cn.seecoder.ai.model.bo.TrainParamsBO
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.RandomForestClassifier

// 随机森林分类算法具体实现类
class RandomForestClassification extends BaseClassification{

  //覆盖父类的模板方法，返回随机森林分类算法的PipelineStage
  // 接收一个TrainParamsBO对象，包含训练参数
  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {
    val rf = new RandomForestClassifier()
      .setFeaturesCol("indexedFeatures")
      .setLabelCol("label")
      .setSeed(trainParams.seed)
      .setMaxDepth(trainParams.maxDepth)
      .setMaxBins(trainParams.maxBins)
      .setMinInstancesPerNode(trainParams.minInstancesPerNode)
      .setMinInfoGain(trainParams.minInfoGain)
      .setNumTrees(trainParams.numTrees)
    rf
  }
}
